#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
创建附件1.csv样本数据
"""

import os
import pandas as pd
import glob
import numpy as np
import datetime

def main():
    print("开始创建附件1.csv样本数据...")
    
    # 获取路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    parent_dir = os.path.dirname(current_dir)
    file_dir = os.path.join(parent_dir, "file")
    output_file = os.path.join(file_dir, "附件1.csv")
    
    # 查找志愿者数据文件
    pattern = os.path.join(file_dir, "P*.csv")
    volunteer_files = glob.glob(pattern)
    
    if not volunteer_files:
        print(f"未找到志愿者数据文件，将创建模拟数据...")
        # 创建模拟数据
        return create_synthetic_data(file_dir, output_file)
    
    print(f"找到 {len(volunteer_files)} 个志愿者数据文件")
    
    # 读取并合并数据 (为了速度，每个文件只读取前1000行)
    all_data = []
    for file in volunteer_files:
        try:
            # 从文件名获取志愿者ID
            volunteer_id = os.path.basename(file).split(".")[0]
            print(f"处理 {volunteer_id} 的数据...")
            
            # 读取数据
            df = pd.read_csv(file, nrows=1000)
            
            # 如果数据中没有pid列，添加一个
            if 'pid' not in df.columns:
                df['pid'] = volunteer_id
                
            # 如果没有timestamp列，创建一个
            if 'timestamp' not in df.columns:
                # 创建一个从2023-01-01开始的时间序列，每条记录间隔10秒
                start_time = datetime.datetime(2023, 1, 1)
                timestamps = [start_time + datetime.timedelta(seconds=i*10) for i in range(len(df))]
                df['timestamp'] = timestamps
                
            # 如果没有MET列，创建一个模拟的MET值
            if 'MET' not in df.columns:
                if all(col in df.columns for col in ['acc_x', 'acc_y', 'acc_z']):
                    print(f"为 {volunteer_id} 计算模拟MET值...")
                    # 用加速度计算简单的MET估算值
                    magnitude = (df['acc_x']**2 + df['acc_y']**2 + df['acc_z']**2)**0.5
                    # 将magnitude映射到0.5-8的MET范围，确保有各种活动类型
                    df['MET'] = 0.5 + 7.5 * (magnitude - magnitude.min()) / (magnitude.max() - magnitude.min() + 0.001)
                    
                    # 增加一些多样性：修改一些数据点以确保各种活动都有
                    # 确保有睡眠 (MET < 1)
                    sleep_indices = np.random.choice(df.index, size=int(len(df)*0.2), replace=False)
                    df.loc[sleep_indices, 'MET'] = np.random.uniform(0.5, 0.99, size=len(sleep_indices))
                    
                    # 确保有静态活动 (1 <= MET < 1.6)
                    static_indices = np.random.choice(df.index, size=int(len(df)*0.15), replace=False)
                    df.loc[static_indices, 'MET'] = np.random.uniform(1.0, 1.59, size=len(static_indices))
                    
                    # 确保有低强度运动 (1.6 <= MET < 3.0)
                    low_indices = np.random.choice(df.index, size=int(len(df)*0.25), replace=False)
                    df.loc[low_indices, 'MET'] = np.random.uniform(1.6, 2.99, size=len(low_indices))
                    
                    # 确保有中等强度运动 (3.0 <= MET < 6.0)
                    medium_indices = np.random.choice(df.index, size=int(len(df)*0.25), replace=False)
                    df.loc[medium_indices, 'MET'] = np.random.uniform(3.0, 5.99, size=len(medium_indices))
                    
                    # 确保有高强度运动 (MET >= 6.0)
                    high_indices = np.random.choice(df.index, size=int(len(df)*0.15), replace=False)
                    df.loc[high_indices, 'MET'] = np.random.uniform(6.0, 8.0, size=len(high_indices))
                else:
                    print(f"为 {volunteer_id} 创建随机MET值...")
                    # 创建多样化的MET值分布
                    met_values = []
                    met_values.extend(np.random.uniform(0.5, 0.99, size=int(len(df)*0.2)))  # 睡眠
                    met_values.extend(np.random.uniform(1.0, 1.59, size=int(len(df)*0.15))) # 静态活动
                    met_values.extend(np.random.uniform(1.6, 2.99, size=int(len(df)*0.25))) # 低强度
                    met_values.extend(np.random.uniform(3.0, 5.99, size=int(len(df)*0.25))) # 中等强度
                    met_values.extend(np.random.uniform(6.0, 8.0, size=int(len(df)*0.15)))  # 高强度
                    
                    # 如果生成的值不够，用中等强度的值填充
                    if len(met_values) < len(df):
                        met_values.extend(np.random.uniform(3.0, 5.99, size=len(df)-len(met_values)))
                    # 如果生成的值过多，截断
                    met_values = met_values[:len(df)]
                    # 打乱顺序
                    np.random.shuffle(met_values)
                    df['MET'] = met_values
            
            all_data.append(df)
            print(f"{volunteer_id} 数据处理完成，读取了 {len(df)} 行")
            
        except Exception as e:
            print(f"处理 {file} 时出错: {e}")
    
    if not all_data:
        print("错误: 没有成功处理任何数据")
        return create_synthetic_data(file_dir, output_file)
    
    # 合并所有数据
    merged_data = pd.concat(all_data, ignore_index=True)
    print(f"合并了 {len(merged_data)} 行数据")
    
    # 保存为附件1.csv
    try:
        merged_data.to_csv(output_file, index=False, encoding='utf-8')
        print(f"样本数据已保存至: {output_file}")
    except Exception as e:
        print(f"保存数据时出错: {e}")

def create_synthetic_data(file_dir, output_file):
    """创建完全模拟的数据"""
    print("创建模拟志愿者数据...")
    
    # 创建志愿者ID列表
    volunteer_ids = [f"P{i:03d}" for i in range(1, 6)]
    
    # 为每个志愿者创建数据
    all_data = []
    for volunteer_id in volunteer_ids:
        print(f"为志愿者 {volunteer_id} 创建模拟数据...")
        
        # 创建随机的记录数量
        n_records = np.random.randint(2000, 5000)
        
        # 创建带时间戳的数据框架
        start_time = datetime.datetime(2023, 1, 1)
        timestamps = [start_time + datetime.timedelta(seconds=i*10) for i in range(n_records)]
        
        # 创建随机加速度数据
        acc_x = np.random.normal(0, 1, size=n_records)
        acc_y = np.random.normal(0, 1, size=n_records)
        acc_z = np.random.normal(0, 1, size=n_records)
        
        # 创建随机陀螺仪数据
        gyro_x = np.random.normal(0, 0.5, size=n_records)
        gyro_y = np.random.normal(0, 0.5, size=n_records)
        gyro_z = np.random.normal(0, 0.5, size=n_records)
        
        # 创建MET值，包含不同活动类型
        met_values = []
        met_values.extend(np.random.uniform(0.5, 0.99, size=int(n_records*0.2)))  # 睡眠
        met_values.extend(np.random.uniform(1.0, 1.59, size=int(n_records*0.15))) # 静态活动
        met_values.extend(np.random.uniform(1.6, 2.99, size=int(n_records*0.25))) # 低强度
        met_values.extend(np.random.uniform(3.0, 5.99, size=int(n_records*0.25))) # 中等强度
        met_values.extend(np.random.uniform(6.0, 8.0, size=int(n_records*0.15)))  # 高强度
        
        # 确保长度正确
        if len(met_values) < n_records:
            met_values.extend(np.random.uniform(3.0, 5.99, size=n_records-len(met_values)))
        met_values = met_values[:n_records]
        
        # 打乱顺序
        np.random.shuffle(met_values)
        
        # 创建数据框
        df = pd.DataFrame({
            'pid': [volunteer_id] * n_records,
            'timestamp': timestamps,
            'acc_x': acc_x,
            'acc_y': acc_y,
            'acc_z': acc_z,
            'gyro_x': gyro_x,
            'gyro_y': gyro_y,
            'gyro_z': gyro_z,
            'MET': met_values
        })
        
        all_data.append(df)
        print(f"为志愿者 {volunteer_id} 创建了 {n_records} 条记录")
    
    # 合并所有数据
    merged_data = pd.concat(all_data, ignore_index=True)
    print(f"共创建了 {len(merged_data)} 条模拟数据记录")
    
    # 保存为附件1.csv
    try:
        merged_data.to_csv(output_file, index=False, encoding='utf-8')
        print(f"模拟数据已保存至: {output_file}")
        return True
    except Exception as e:
        print(f"保存数据时出错: {e}")
        return False

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"程序运行出错: {e}")
    
    print("\n按回车键退出...")
    input() 